# coding:utf-8
# 导入必备工具包
import torch
from rich import print


# ======================== MLM Loss 计算函数 ========================
def mlm_loss(logits, mask_positions, sub_mask_labels,
             cross_entropy_criterion, device):
    """
    计算指定位置的 mask token 的 output 与 label 之间的 cross entropy loss。

    Args:
        logits (torch.tensor): 模型原始输出 -> (batch, seq_len, vocab_size)
        mask_positions (torch.tensor): mask token 的位置 -> (batch, mask_label_num)
        sub_mask_labels (list): mask token 的 sub label, 由于每个 label 的 sub_label 数目不同，
                                所以这里是个变长的 list, e.g. -> [
                                    [[2398, 3352]],
                                    [[2398, 3352], [3819, 3861]]
                                ]
        cross_entropy_criterion (CrossEntropyLoss): CE Loss 计算器
        device (str): cpu 还是 gpu

    Returns:
        torch.tensor: CE Loss
    """
    batch_size, seq_len, vocab_size = logits.size()  # 获取模型输出的形状参数
    loss = None  # 初始化 loss 变量

    for single_value in zip(logits, sub_mask_labels, mask_positions):  # 遍历每个样本及其对应的 mask 位置和标签
        single_logits = single_value[0]  # 当前样本的 logits -> (seq_len, vocab_size)
        single_sub_mask_labels = single_value[1]  # 当前样本的 sub_mask_labels
        single_mask_positions = single_value[2]  # 当前样本的 mask_positions -> (mask_label_num)

        single_mask_logits = single_logits[single_mask_positions]  # 提取 mask 位置的 logits -> (mask_label_num, vocab_size)
        single_mask_logits = single_mask_logits.repeat(len(single_sub_mask_labels), 1,
                                                       1)  # 复制 logits 以匹配 sub_mask_labels 的长度
        single_mask_logits = single_mask_logits.reshape(-1,
                                                        vocab_size)  # 将 logits 展平为 (sub_label_num * mask_label_num, vocab_size)

        single_sub_mask_labels = torch.LongTensor(single_sub_mask_labels).to(
            device)  # 转换 sub_mask_labels 为 Tensor 并移动到指定设备
        single_sub_mask_labels = single_sub_mask_labels.reshape(-1,
                                                                1).squeeze()  # 将 sub_mask_labels 展平为 (sub_label_num * mask_label_num)

        cur_loss = cross_entropy_criterion(single_mask_logits, single_sub_mask_labels)  # 计算当前样本的交叉熵损失
        cur_loss = cur_loss / len(single_sub_mask_labels)  # 对当前样本的损失进行归一化

        if not loss:  # 如果 loss 为空，初始化为当前样本的损失
            loss = cur_loss
        else:  # 否则将当前样本的损失累加到总损失中
            loss += cur_loss

    loss = loss / batch_size  # 对总损失进行批量归一化 -> (1,)
    return loss  # 返回最终的 CE Loss


# ======================== Logits 转 Token ID 函数 ========================
def convert_logits_to_ids(
        logits: torch.tensor,
        mask_positions: torch.tensor):
    """
    输入 Language Model 的词表概率分布（LMModel 的 logits），将 mask_position 位置的
    token logits 转换为 token 的 id。

    Args:
        logits (torch.tensor): model output -> (batch, seq_len, vocab_size)
        mask_positions (torch.tensor): mask token 的位置 -> (batch, mask_label_num)

    Returns:
        torch.LongTensor: 对应 mask position 上最大概率的推理 token -> (batch, mask_label_num)
    """
    label_length = mask_positions.size()[1]  # 获取每个样本的 mask 标签数量
    batch_size, seq_len, vocab_size = logits.size()  # 获取 logits 的形状参数

    mask_positions_after_reshaped = []  # 初始化存储展平后的 mask 位置索引
    for batch, mask_pos in enumerate(mask_positions.detach().cpu().numpy().tolist()):  # 遍历每个样本的 mask 位置
        for pos in mask_pos:  # 遍历单个样本中的每个 mask 位置
            mask_positions_after_reshaped.append(batch * seq_len + pos)  # 计算全局索引并存储

    logits = logits.reshape(batch_size * seq_len, -1)  # 将 logits 展平为 (batch_size * seq_len, vocab_size)
    mask_logits = logits[mask_positions_after_reshaped]  # 提取对应 mask 位置的 logits -> (batch * label_num, vocab_size)
    predict_tokens = mask_logits.argmax(dim=-1)  # 找到每个 mask 位置上概率最大的 token id -> (batch * label_num)
    predict_tokens = predict_tokens.reshape(-1, label_length)  # 将预测结果重塑为 (batch, label_num)

    return predict_tokens  # 返回预测的 token id


# ======================== 主程序测试部分 ========================
if __name__ == '__main__':
    # 测试数据
    logits = torch.randn(2, 20, 21193)  # 随机生成一个模拟 logits 张量 -> (batch=2, seq_len=20, vocab_size=21193)
    mask_positions = torch.LongTensor([  # 定义 mask 位置张量 -> (batch=2, mask_label_num=2)
        [5, 6],
        [5, 6],
    ])

    # 测试 convert_logits_to_ids 函数
    predict_tokens = convert_logits_to_ids(logits, mask_positions)  # 调用函数将 logits 转换为预测的 token id
    print("Predicted Tokens:")
    print(predict_tokens)  # 打印预测结果

    # 测试 mlm_loss 函数
    sub_mask_labels = [  # 定义 sub_mask_labels
        [[2398, 3352]],
        [[2398, 3352], [3819, 3861]]
    ]
    cross_entropy_criterion = torch.nn.CrossEntropyLoss()  # 定义交叉熵损失计算器
    device = "cpu"  # 使用 CPU 设备

    loss = mlm_loss(logits, mask_positions, sub_mask_labels, cross_entropy_criterion, device)  # 计算 MLM 损失
    print("\nMLM Loss:")
    print(loss.item())  # 打印损失值
# 这段代码整合了两个核心功能：
# MLM 损失计算：适用于多标签场景，支持动态长度的子标签。
# Logits 转 Token ID：推断出每个 mask 位置上的最可能 token id。
# 整体设计灵活且高效，适用于基于 Prompt 的自然语言处理任务（如 MLM、文本分类等）。代码充分利用了 PyTorch 的张量操作和 Hugging Face 的工具链，具有良好的扩展性和通用性。